﻿using CsvHelper.Configuration;
using CsvHelper;
using Microsoft.ML.OnnxRuntime.Tensors;
using Microsoft.ML.OnnxRuntime;
using System;
using System.Collections.Generic;
using System.Globalization;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace DataWorks_Tools.LSTMDetails
{
    public static class OnnxTools
    {

        /// <summary>
        /// 根据模型名称来选取对应的scalers
        /// </summary>
        /// <param name="onnxfile"></param>
        /// <param name="scalers"></param>
        /// <returns></returns>
        /// <exception cref="DirectoryNotFoundException"></exception>
        /// <exception cref="FileNotFoundException"></exception>
        public static FileInfo[] SelectScaler(FileInfo onnxfile,FileInfo[] scalers)
        {
            //获取模型对应的wft名称
            var wftname= onnxfile.Name.Replace("_model.onnx", "");
            // 在 scalers 数组中查找对应的两个文件
            var xScaler = scalers.FirstOrDefault(f => f.Name == $"{wftname}_x_scaler.csv");
            var yScaler = scalers.FirstOrDefault(f => f.Name == $"{wftname}_y_scaler.csv");

            if (xScaler == null || yScaler == null)
                throw new FileNotFoundException($"未找到匹配 {wftname}_x_scaler.csv 或 {wftname}_y_scaler.csv 的文件");

            return new FileInfo[] { xScaler, yScaler }; // x 在前，y 在后
        }
        public static FileInfo SelectYScaler(FileInfo onnxfile, FileInfo[] scalers)
        {
            //获取模型对应的wft名称
            var wftname = onnxfile.Name.Replace("_model.onnx", "");
            // 在 scalers 数组中查找对应的两个文件
            var yScaler = scalers.FirstOrDefault(f => f.Name == $"{wftname}_y_scaler.csv");

            if ( yScaler == null)
                throw new FileNotFoundException($"未找到匹配 {wftname}_y_scaler.csv 的文件");

            return yScaler; // x 在前，y 在后
        }




        /// <summary>
        /// 将一个包含连续特征向量的 List<float[]> 按照时间窗口（50 个时间步）组织成符合 ONNX LSTM 模型输入格式（[1, 50, 16]）的张量 DenseTensor<float> 列表。
        /// </summary>
        /// <param name="normData"></param>
        /// <returns></returns>
        public static List<DenseTensor<float>> CreateLstmInputTensors(List<float[]> normData)
        {
            int featureSize = normData[0].Length;
            int timeSteps = 50;
            int totalSamples = normData.Count - timeSteps + 1;

            var tensors = new List<DenseTensor<float>>();

            for (int i = 0; i < totalSamples; i++)
            {
                var tensor = new DenseTensor<float>(new[] { 1, timeSteps, featureSize });

                for (int t = 0; t < timeSteps; t++)
                {
                    float[] stepData = normData[i + t];
                    for (int f = 0; f < featureSize; f++)
                    {
                        tensor[0, t, f] = stepData[f];
                    }
                }

                tensors.Add(tensor);
            }

            return tensors;
        }

        /// <summary>
        /// 使用 ONNX 模型批量预测 LSTM 输入数据。
        /// </summary>
        /// <param name="session">ONNX 运行会话</param>
        /// <param name="inputTensors">模型输入张量列表（每个 shape 为 [1, timeSteps, featureSize]）</param>
        /// <param name="inputName">模型输入名称，默认是 "input"</param>
        /// <returns>预测结果列表，每个结果为 float[]</returns>
        public static List<float[]> PredictBatch(InferenceSession session, List<DenseTensor<float>> inputTensors, string inputName = "input")
        {
            var allPredictions = new List<float[]>();

            foreach (var tensor in inputTensors)
            {
                var inputs = new List<NamedOnnxValue>
            {
                NamedOnnxValue.CreateFromTensor(inputName, tensor)
            };

                using IDisposableReadOnlyCollection<DisposableNamedOnnxValue> results = session.Run(inputs);
                var resultTensor = results.First().AsEnumerable<float>().ToArray();

                allPredictions.Add(resultTensor);
            }

            return allPredictions;
        }


        /// <summary>
        /// 将 List<float[]> 扁平化为一个 List<float>。
        /// </summary>
        /// <param name="list">要合并的 List<float[]></param>
        /// <returns>合并后的一维 List<float></returns>
        public static List<float> FlattenToList(List<float[]> list)
        {
            if (list == null || list.Count == 0)
                return new List<float>();

            var result = new List<float>(list.Sum(arr => arr.Length));
            foreach (var arr in list)
            {
                result.AddRange(arr);
            }

            return result;
        }
        /// <summary>
        /// 只负责根据 data 的第一个元素的长度，生成 50 个 [0, 0, ..., 0] 的 float[] 元素组成的 List<float[]>
        /// </summary>
        /// <param name="data"></param>
        /// <param name="paddingCount"></param>
        /// <returns></returns>
        /// <exception cref="ArgumentException"></exception>
        public static List<float[]> CreateZeroPadding(List<float[]> data, int paddingCount = 50-1)
        {
            if (data == null || data.Count == 0)
                throw new ArgumentException("data 不能为空且必须包含至少一个元素");
            int featureLength = data[0].Length;
            var padding = new List<float[]>(paddingCount);
            for (int i = 0; i < paddingCount; i++)
            {
                padding.Add(new float[featureLength]); // 默认初始化为全 0
            }
            return padding;
        }

        /// <summary>
        /// 利用csvhelper来保存csv
        /// </summary>
        /// <param name="data"></param>
        /// <param name="headers"></param>
        /// <param name="path"></param>
        /// <returns></returns>
        public static async Task SaveCSVfromArrayData(List<List<float>> data, List<string> headers, string path)
        {
            var directory = Path.GetDirectoryName(path);
            if (!string.IsNullOrEmpty(directory) && !Directory.Exists(directory))
            {
                Directory.CreateDirectory(directory);
            }

            using (var writer = new StreamWriter(path, false, Encoding.UTF8))
            using (var csvWriter = new CsvWriter(writer, new CsvConfiguration(CultureInfo.InvariantCulture)))
            {
                // 写入表头
                foreach (var column in headers)
                {
                    csvWriter.WriteField(column);
                }
                await csvWriter.NextRecordAsync();

                int rowCount = data[0].Count;  // 每列的行数
                int colCount = data.Count;     // 列数（= header.Count）

                for (int i = 0; i < rowCount; i++) // 对于每一行
                {
                    for (int j = 0; j < colCount; j++) // 对于每一列
                    {
                        csvWriter.WriteField(data[j][i]); // 第 j 列的第 i 行数据
                    }
                    await csvWriter.NextRecordAsync();
                }
            }
        }

        /// <summary>
        /// 高性能纯字符串拼接版本的 SaveCSVfromArrayData 方法，适用于：数据量很大（上万行以上）；headers 和 data 中的字段不会包含 逗号 ,、引号 " 或换行符；不依赖 CsvHelper（更轻量，适合生产环境中数据处理导出）。
        /// </summary>
        /// <param name="data"></param>
        /// <param name="headers"></param>
        /// <param name="path"></param>
        public static void SaveCSVFast(List<List<float>> data, List<string> headers, string path)
        {
            var directory = Path.GetDirectoryName(path);
            if (!string.IsNullOrEmpty(directory) && !Directory.Exists(directory))
            {
                Directory.CreateDirectory(directory);
            }

            using var writer = new StreamWriter(path, false, Encoding.UTF8);

            // 写入表头
            writer.WriteLine(string.Join(",", headers));

            // 写入每一行数据
            foreach (var row in data)
            {
                var line = string.Join(",", row.Select(cell => cell.ToString("G17"))); // 保证精度
                writer.WriteLine(line);
            }

            writer.Flush();
        }
    
    }
}
